# encoding: UTF-8
"""
条件库
"""
from base.baseCondiction import CondictionTemplate
import numpy as np


class MACondiction(CondictionTemplate):
    """MA 条件指标"""
    def __init__(self):
        """条件初始化"""
        super(MACondiction, self).__init__()

        self.hisMA30 = []
        self.hisMA35 = []
        self.hisMA40 = []
class MACDCondiction(CondictionTemplate):
    """MACD 条件指标"""
    def __init__(self):

        super(MACDCondiction, self).__init__()

        self.hisShort = [] #短期EMA均线
        self.hisLong = [] #长期EMA均线
    def caculate(self):
        """计算"""

        if len(self.hisShort) == 0:
            self.hisShort.append(self.hisBar[-1].close)
            self.hisLong.append(self.hisBar[-1].close)
            return
        shortEMA = (2 * self.hisBar[-1].close + 11 * self.hisShort[-1]) / 13.0
        longEMA = (2 * self.hisBar[-1].close + 25 * self.hisLong[-1]) / 27.0

        self.hisShort.append(shortEMA)
        self.hisLong.append(longEMA)
            
    def buy(self):
        """买入条件"""
        signal = False
        if len(self.hisShort)>=2:
            #短期EMA向上突破长期EMA
            signal = self.hisShort[-2] < self.hisLong[-2] and self.hisShort[-1] > self.hisLong[-1]
        return signal

    def sell(self):
        """买入平仓条件"""
        signal = False

        if len(self.hisShort)>=2:
            #收盘价向下突破短期EMA
            pS = self.hisBar[-2].close > self.hisShort[-2] and self.hisBar[-1].close < self.hisShort[-1]
            #短期EMA向下突破长期EMA
            eS = self.hisShort[-2] > self.hisLong[-2] and self.hisShort[-1] < self.hisLong[-1]
            signal = pS or eS
        return signal

    def short(self):
        """卖空条件"""
        signal = False
        if len(self.hisShort)>=2:
            #短期EMA向下突破长期EMA
            signal = self.hisShort[-2] > self.hisLong[-2] and self.hisShort[-1] < self.hisLong[-1]
        return signal
    def cover(self):
        """卖出平仓条件"""
        signal = False
        if len(self.hisShort)>=2:
            #收盘价向上突破EMA
            pS = self.hisBar[-2].close  < self.hisShort[-2]
            #短期EMA向上突破长期EMA
            eS = self.hisShort[-2] < self.hisLong[-2] and self.hisShort[-1] > self.hisLong[-1]
            signal = pS or eS
        return signal

    def getSignalStr(self):
        if len(self.hisShort) >0:
            return "datetime=%s short=%f , long=%f" % (self.hisBar[-1].datetime, self.hisShort[-1], self.hisLong[-1])
        return ""


class MACDCondiction2(CondictionTemplate):
    """MACD 条件指标"""
    """
    指标: 1.判断MACD的绿变红，买入开仓，并保留10周期
          2.判断MACD的红变绿，卖出开仓，并保留10周期
    """
    def __init__(self):

        super(MACDCondiction2, self).__init__()

        self.hisShort = [] #短期EMA均线
        self.hisLong = [] #长期EMA均线
        self.hisDiff = [] #DIFF
        self.hisDea = [] #DEA
        self.hisMBAR = [] #红绿柱

    def caculate(self):
        """计算"""

        if len(self.hisShort) == 0:
            self.hisShort.append(self.hisBar[-1].close)
            self.hisLong.append(self.hisBar[-1].close)
            return
        shortEMA = (2 * self.hisBar[-1].close + 11 * self.hisShort[-1]) / 13.0
        longEMA = (2 * self.hisBar[-1].close + 25 * self.hisLong[-1]) / 27.0
        DIFF = shortEMA - longEMA
 
        if (len(self.hisDiff) == 0):
            self.hisDiff.append(DIFF)
            return
        if len(self.hisDea) >0: 
            DEA = 8 * self.hisDea[-1] / 10.0 + 2 * DIFF / 10.0
        else:
            DEA = 2 * DIFF / 10.0

        MBAR = DIFF - DEA

        self.hisShort.append(shortEMA)
        self.hisLong.append(longEMA)
        self.hisDiff.append(DIFF)
        self.hisDea.append(DEA)
        self.hisMBAR.append(MBAR)

            
    def buy(self):
        """买入条件"""
        signal = False
        if len(self.hisMBAR)>=2:
            #BAR柱绿变红买入
            signal = self.hisMBAR[-2] < 0 and self.hisMBAR[-1] > 0
        return signal

    def sell(self):
        """卖出平仓条件"""
        signal = False
        """
        持仓10个BAR
        """
        
        return signal

    def short(self):
        """卖空条件"""
        signal = False
        #if len(self.hisMBAR)>=2:
        #    #BAR柱绿变红买入
        #    signal = self.hisMBAR[-2] > 0 and self.hisMBAR[-1] < 0
        return signal
    def cover(self):
        """买入平仓条件"""
        signal = False
        """
        持仓10个BAR
        """
        return signal

    def getSignalStr(self):
        if len(self.hisShort) >0:
            return "datetime=%s short=%f ,long=%f, DIFF=%f, DEA=%f, MBAR=%f" % (self.hisBar[-1].datetime, self.hisShort[-1], self.hisLong[-1],\
                            self.hisDiff[-1], self.hisDea[-1], self.hisMBAR[-1])
        return ""

class MALineTradeCondiction(CondictionTemplate):
    """MA线上交易开多, MA线开空"""
    def __init__(self):
        """条件初始化"""
        super(MALineTradeCondiction, self).__init__()

        self.hisMA = []

    def caculate(self):
        """计算"""
        if len(self.hisBar) < self.param['MALine-PERIOD']:
            return
        
        MA =  sum(self.hisBarClose[-self.param['MALine-PERIOD']:]) / self.param['MALine-PERIOD']
        self.hisMA.append(MA)

    def buy(self):
        if len(self.hisMA) < 1:
            return False
        if self.hisBarClose >= self.hisMA[-1]:
            return True
        return False
    def sell(self):
        if len(self.hisMA) < 1:
            return False
        if self.hisBarClose < self.hisMA[-1]:
            return True
        return False

    def short(self):
        return self.sell()
    
    def cover(self):
        return self.buy()
    
    def getSignalStr(self):
        if len(self.hisMA) >0:
            return "datetime=%s MA=%f" % (self.hisBar[-1], self.hisMA[-1])
        return ""
class STOCondiction(CondictionTemplate):
    """STO 条件指标"""
    def __init__(self):
        super(STOCondiction, self).__init__()
        
        self.hisK = []
        self.hisD = []
        self.hisDS = []
    
    def caculate(self):
        """计算"""
        if len(self.hisBar) < self.param['STO-PERIOD']:
            return
        
        Low5 = min(self.hisBarLow[-5:])
        High5 = max(self.hisBarHigh[-5:])
        if Low5 == High5:
            return 
        K = 100 * ((self.hisBar[-1].close - Low5) / (High5 - Low5))
        self.hisK.append(K)

        if len(self.hisK) < self.param['STO-KPERIOD']:
            return
        D = (sum(self.hisK[-self.param['STO-DPERIOD']:]) / float(self.param['STO-DPERIOD']))
        self.hisD.append(D)
        DS = sum(self.hisD[-self.param['STO-DPERIOD']:]) / float(self.param['STO-DPERIOD'])
        self.hisDS.append(DS)
    def buy(self):

        if len(self.hisDS) <2:
            return False
        if self.hisK[-1] <= self.param['STO-BUY'] and self.hisDS[-1] <= self.param['STO-BUY'] \
            and self.hisK[-2] < self.hisDS[-2] and self.hisK[-1] >= self.hisDS[-1]:
            return True
        return False

    def sell(self):
        if len(self.hisDS) < 2:
            return False
        if (self.hisK[-2] < self.param['STO-SELL'] and self.hisK[-1] >= self.param['STO-SELL'])\
           or (self.hisDS[-2] < self.param['STO-SELL'] and self.hisDS[-1] >= self.param['STO-SELL']):
            return True
        return False
    def short(self):
        if len(self.hisDS) < 2:
            return False
        if self.hisK[-1] >= self.param['STO-SELL'] and self.hisDS[-1] >=self.param['STO-SELL']\
                        and self.hisK[-2] > self.hisDS[-2] and self.hisK[-2] <= self.hisDS[-1]:
            return True
        return False
    def cover(self):
        if len(self.hisDS) < 2:
            return False
        if (self.hisK[-2] > self.param['STO-BUY'] and self.hisK[-1] <= self.param['STO-BUY'])\
            or (self.hisDS[-2] > self.param['STO-BUY'] and self.hisDS[-1] <= self.param['STO-BUY']):
            return True
        return False
    def getSignalStr(self):
        if (len(self.hisK) >0) and (len(self.hisDS) >0):
            return "datetime=%s, K=%s, DS=%s" % (self.hisBar[-1].datetime, self.hisK[-1], self.hisDS[-1])



class RSICondiction(CondictionTemplate):
    """RSI 条件指标"""
    def __init__(self):
        """条件初始化"""
        super(RSICondiction, self).__init__()

        self.RS = 0.0
        self.RSI = 0
        self.TotalUp = 0.0
        self.TotalDown = 0.0

        self.hisRS = []
        self.hisRSI = []

    def caculate(self):
        """计算"""
        if len(self.hisBar) < self.param['RSI-PERIOD']:
            return 
        prices = np.array(self.hisBarClose[-self.param['RSI-PERIOD']:]) - \
                           np.array(self.hisBarOpen[-self.param['RSI-PERIOD']:])
        self.TotalUp = sum([i for i in prices if i >= 0])
        self.TotalDown = sum([i for i in prices if i < 0])
        if self.TotalUp == 0 or self.TotalDown == 0:
            return
        self.RS = - self.TotalUp / self.TotalDown
        self.RSI = 100 * (self.RS) / (1 + self.RS)
        self.hisRS.append(self.RS)
        self.hisRSI.append(self.RSI)

    def buy(self):
        if len(self.hisRSI) <2:
            return False

        if self.hisRSI[-2] < self.param['RSI-BUY'] and self.hisRSI[-1] >=self.param['RSI-BUY'] and self.RSI != 0.0:
            return True
        return False
    def sell(self):
        if len(self.hisRSI) <2:
            return False

        if self.hisRSI[-2] < self.param['RSI-SELL'] and self.hisRSI[-1] >= self.param['RSI-SELL'] and self.RSI !=0.0:
            return True
        return False
    def short(self):
        if len(self.hisRSI) <2:
            return False
    
        if self.hisRSI[-2] > self.param['RSI-SELL'] and self.hisRSI[-1] <= self.param['RSI-SELL'] and self.RSI !=0.0:
            return True
        return False
    def cover(self):
        if len(self.hisRSI) <2:
            return False

        if self.hisRSI[-2] > self.param['RSI-BUY']  and self.hisRSI[-1] <= self.param['RSI-BUY'] and self.RSI !=0.0:
            return True
        return False

    def getSignalStr(self):
        return "datetime=%s, RS=%f, RSI=%d " % (self.hisBar[-1].datetime, self.RS, self.RSI)
        


